{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 225,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "\n",
    "from torch.nn.utils import clip_grad_norm_\n",
    "import random\n",
    "import math\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from collections import deque, namedtuple\n",
    "import time\n",
    "import gym\n",
    "import copy\n",
    "\n",
    "from torch.distributions import MultivariateNormal, Normal\n",
    "#torch.autograd.set_detect_anomaly(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 252,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NAF(nn.Module):\n",
    "    def __init__(self, state_size, action_size,layer_size, seed):\n",
    "        super(NAF, self).__init__()\n",
    "        self.seed = torch.manual_seed(seed)\n",
    "        self.input_shape = state_size\n",
    "        self.action_size = action_size\n",
    "                \n",
    "        self.head_1 = nn.Linear(self.input_shape, layer_size)\n",
    "        self.bn1 = nn.BatchNorm1d(layer_size)\n",
    "        self.ff_1 = nn.Linear(layer_size, layer_size)\n",
    "        self.bn2 = nn.BatchNorm1d(layer_size)\n",
    "        self.action_values = nn.Linear(layer_size, action_size)\n",
    "        self.value = nn.Linear(layer_size, 1)\n",
    "        self.matrix_entries = nn.Linear(layer_size, int(self.action_size*(self.action_size+1)/2))\n",
    "        \n",
    "\n",
    "    \n",
    "    def forward(self, input, action=None):\n",
    "        \"\"\"\n",
    "        \n",
    "        \"\"\"\n",
    "\n",
    "        x = torch.relu(self.head_1(input))\n",
    "        x = torch.relu(self.ff_1(x))\n",
    "        action_value = torch.tanh(self.action_values(x))\n",
    "        entries = torch.tanh(self.matrix_entries(x))\n",
    "        V = self.value(x)\n",
    "        \n",
    "        action_value = action_value.unsqueeze(-1)\n",
    "        \n",
    "        # create lower-triangular matrix\n",
    "        L = torch.zeros((input_.shape[0], self.action_size, self.action_size)).to(device)\n",
    "\n",
    "        # get lower triagular indices\n",
    "        tril_indices = torch.tril_indices(row=self.action_size, col=self.action_size, offset=0)  \n",
    "\n",
    "        # fill matrix with entries\n",
    "        L[:, tril_indices[0], tril_indices[1]] = entries\n",
    "        L.diagonal(dim1=1,dim2=2).exp_()\n",
    "\n",
    "        # calculate state-dependent, positive-definite square matrix\n",
    "        P = L*L.transpose(2, 1)\n",
    "        \n",
    "        Q = None\n",
    "        if action is not None:\n",
    "\n",
    "            # calculate Advantage:\n",
    "            A = (-0.5 * torch.matmul(torch.matmul((action.unsqueeze(-1) - action_value).transpose(2, 1), P), (action.unsqueeze(-1) - action_value))).squeeze(-1)\n",
    "\n",
    "            Q = A + V\n",
    "        \n",
    "        \n",
    "        # add noise to action mu:\n",
    "        dist = MultivariateNormal(action_value.squeeze(-1), torch.inverse(P))\n",
    "        #dist = Normal(action_value.squeeze(-1), 1)\n",
    "        action = dist.sample()\n",
    "        action = torch.clamp(action, min=-1, max=1)\n",
    "        #action = action_value.squeeze(-1)\n",
    "        \n",
    "        return action, Q, V\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 253,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    \"\"\"Fixed-size buffer to store experience tuples.\"\"\"\n",
    "\n",
    "    def __init__(self, buffer_size, batch_size, device, seed, gamma):\n",
    "        \"\"\"Initialize a ReplayBuffer object.\n",
    "        Params\n",
    "        ======\n",
    "            buffer_size (int): maximum size of buffer\n",
    "            batch_size (int): size of each training batch\n",
    "            seed (int): random seed\n",
    "        \"\"\"\n",
    "        self.device = device\n",
    "        self.memory = deque(maxlen=buffer_size)  \n",
    "        self.batch_size = batch_size\n",
    "        self.experience = namedtuple(\"Experience\", field_names=[\"state\", \"action\", \"reward\", \"next_state\", \"done\"])\n",
    "        self.seed = random.seed(seed)\n",
    "        self.gamma = gamma\n",
    "\n",
    "        self.n_step_buffer = deque(maxlen=1)\n",
    "    \n",
    "    def add(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Add a new experience to memory.\"\"\"\n",
    "\n",
    "        self.n_step_buffer.append((state, action, reward, next_state, done))\n",
    "        if len(self.n_step_buffer) == 1:\n",
    "            state, action, reward, next_state, done = self.calc_multistep_return()\n",
    "\n",
    "            e = self.experience(state, action, reward, next_state, done)\n",
    "            self.memory.append(e)\n",
    "    \n",
    "    def calc_multistep_return(self):\n",
    "        Return = 0\n",
    "        for idx in range(1):\n",
    "            Return += self.gamma**idx * self.n_step_buffer[idx][2]\n",
    "        \n",
    "        return self.n_step_buffer[0][0], self.n_step_buffer[0][1], Return, self.n_step_buffer[-1][3], self.n_step_buffer[-1][4]\n",
    "        \n",
    "    \n",
    "    \n",
    "    def sample(self):\n",
    "        \"\"\"Randomly sample a batch of experiences from memory.\"\"\"\n",
    "        experiences = random.sample(self.memory, k=self.batch_size)\n",
    "\n",
    "        states = torch.from_numpy(np.stack([e.state for e in experiences if e is not None])).float().to(self.device)\n",
    "        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).long().to(self.device)\n",
    "        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(self.device)\n",
    "        next_states = torch.from_numpy(np.stack([e.next_state for e in experiences if e is not None])).float().to(self.device)\n",
    "        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(self.device)\n",
    "  \n",
    "        return (states, actions, rewards, next_states, dones)\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Return the current size of internal memory.\"\"\"\n",
    "        return len(self.memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 254,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OUNoise:\n",
    "    \"\"\"Ornstein-Uhlenbeck process.\"\"\"\n",
    "\n",
    "    def __init__(self, size, seed, mu=0., theta=0.15, sigma=0.2):\n",
    "        \"\"\"Initialize parameters and noise process.\"\"\"\n",
    "        self.mu = mu * np.ones(size)\n",
    "        self.theta = theta\n",
    "        self.sigma = sigma\n",
    "        self.seed = random.seed(seed)\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"Reset the internal state (= noise) to mean (mu).\"\"\"\n",
    "        self.state = copy.copy(self.mu)\n",
    "\n",
    "    def sample(self):\n",
    "        \"\"\"Update internal state and return it as a noise sample.\"\"\"\n",
    "        x = self.state\n",
    "        dx = self.theta * (self.mu - x) + self.sigma * np.array([random.random() for i in range(len(x))])\n",
    "        self.state = x + dx\n",
    "        return self.state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 255,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQN_Agent():\n",
    "    \"\"\"Interacts with and learns from the environment.\"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 state_size,\n",
    "                 action_size,\n",
    "                 layer_size,\n",
    "                 BATCH_SIZE,\n",
    "                 BUFFER_SIZE,\n",
    "                 LR,\n",
    "                 TAU,\n",
    "                 GAMMA,\n",
    "                 UPDATE_EVERY,\n",
    "                 NUPDATES,\n",
    "                 device,\n",
    "                 seed):\n",
    "        \"\"\"Initialize an Agent object.\n",
    "        \n",
    "        Params\n",
    "        ======\n",
    "            state_size (int): dimension of each state\n",
    "            action_size (int): dimension of each action\n",
    "            Network (str): dqn network type\n",
    "            layer_size (int): size of the hidden layer\n",
    "            BATCH_SIZE (int): size of the training batch\n",
    "            BUFFER_SIZE (int): size of the replay memory\n",
    "            LR (float): learning rate\n",
    "            TAU (float): tau for soft updating the network weights\n",
    "            GAMMA (float): discount factor\n",
    "            UPDATE_EVERY (int): update frequency\n",
    "            device (str): device that is used for the compute\n",
    "            seed (int): random seed\n",
    "        \"\"\"\n",
    "        self.state_size = state_size\n",
    "        self.action_size = action_size\n",
    "        self.seed = random.seed(seed)\n",
    "        self.device = device\n",
    "        self.TAU = TAU\n",
    "        self.GAMMA = GAMMA\n",
    "        self.UPDATE_EVERY = UPDATE_EVERY\n",
    "        self.NUPDATES = NUPDATES\n",
    "        self.BATCH_SIZE = BATCH_SIZE\n",
    "        self.Q_updates = 0\n",
    "\n",
    "\n",
    "        self.action_step = 4\n",
    "        self.last_action = None\n",
    "\n",
    "        # Q-Network\n",
    "        self.qnetwork_local = NAF(state_size, action_size,layer_size, seed).to(device)\n",
    "        self.qnetwork_target = NAF(state_size, action_size,layer_size, seed).to(device)\n",
    "\n",
    "        self.optimizer = optim.Adam(self.qnetwork_local.parameters(), lr=LR)\n",
    "        print(self.qnetwork_local)\n",
    "        \n",
    "        # Replay memory\n",
    "        self.memory = ReplayBuffer(BUFFER_SIZE, BATCH_SIZE, self.device, seed, self.GAMMA)\n",
    "        \n",
    "        # Initialize time step (for updating every UPDATE_EVERY steps)\n",
    "        self.t_step = 0\n",
    "        \n",
    "        # Noise process\n",
    "        self.noise = OUNoise(action_size, seed)\n",
    "    \n",
    "    def step(self, state, action, reward, next_state, done, writer):\n",
    "        # Save experience in replay memory\n",
    "        self.memory.add(state, action, reward, next_state, done)\n",
    "        \n",
    "        # Learn every UPDATE_EVERY time steps.\n",
    "        self.t_step = (self.t_step + 1) % self.UPDATE_EVERY\n",
    "        if self.t_step == 0:\n",
    "            # If enough samples are available in memory, get random subset and learn\n",
    "            if len(self.memory) > self.BATCH_SIZE:\n",
    "                Q_losses = []\n",
    "                for _ in range(self.NUPDATES):\n",
    "                    experiences = self.memory.sample()\n",
    "                    loss = self.learn(experiences)\n",
    "                    self.Q_updates += 1\n",
    "                    Q_losses.append(loss)\n",
    "                writer.add_scalar(\"Q_loss\", np.mean(Q_losses), self.Q_updates)\n",
    "\n",
    "    def act(self, state):\n",
    "        \"\"\"Calculating the action\n",
    "        \n",
    "        Params\n",
    "        ======\n",
    "            state (array_like): current state\n",
    "            \n",
    "        \"\"\"\n",
    "\n",
    "        state = torch.from_numpy(state).float().to(self.device)\n",
    "\n",
    "        self.qnetwork_local.eval()\n",
    "        with torch.no_grad():\n",
    "            action, _, _ = self.qnetwork_local(state.unsqueeze(0))\n",
    "            #action = action.cpu().squeeze().numpy() + self.noise.sample()\n",
    "            #action = np.clip(action, -1,1)[0]\n",
    "        self.qnetwork_local.train()\n",
    "        return action.cpu().squeeze().numpy()\n",
    "\n",
    "\n",
    "\n",
    "    def learn(self, experiences):\n",
    "        \"\"\"Update value parameters using given batch of experience tuples.\n",
    "        Params\n",
    "        ======\n",
    "            experiences (Tuple[torch.Tensor]): tuple of (s, a, r, s', done) tuples \n",
    "        \"\"\"\n",
    "        self.optimizer.zero_grad()\n",
    "        states, actions, rewards, next_states, dones = experiences\n",
    "\n",
    "        # get the Value for the next state from target model\n",
    "        with torch.no_grad():\n",
    "            _, _, V_ = self.qnetwork_target(next_states)\n",
    "\n",
    "        # Compute Q targets for current states \n",
    "        V_targets = rewards + (self.GAMMA * V_ * (1 - dones))\n",
    "        \n",
    "        # Get expected Q values from local model\n",
    "        _, Q, _ = self.qnetwork_local(states, actions)\n",
    "\n",
    "        # Compute loss\n",
    "        loss = F.mse_loss(Q, V_targets)\n",
    "        \n",
    "        # Minimize the loss\n",
    "        loss.backward()\n",
    "        clip_grad_norm_(self.qnetwork_local.parameters(),1)\n",
    "        self.optimizer.step()\n",
    "\n",
    "        # ------------------- update target network ------------------- #\n",
    "        self.soft_update(self.qnetwork_local, self.qnetwork_target)\n",
    "        \n",
    "        self.noise.reset()\n",
    "        \n",
    "        return loss.detach().cpu().numpy()            \n",
    "\n",
    "    def soft_update(self, local_model, target_model):\n",
    "        \"\"\"Soft update model parameters.\n",
    "        θ_target = τ*θ_local + (1 - τ)*θ_target\n",
    "        Params\n",
    "        ======\n",
    "            local_model (PyTorch model): weights will be copied from\n",
    "            target_model (PyTorch model): weights will be copied to\n",
    "            tau (float): interpolation parameter \n",
    "        \"\"\"\n",
    "        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):\n",
    "            target_param.data.copy_(self.TAU*local_param.data + (1.0-self.TAU)*target_param.data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using  cpu\n",
      "NAF(\n",
      "  (head_1): Linear(in_features=4, out_features=256, bias=True)\n",
      "  (bn1): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (ff_1): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (bn2): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (action_values): Linear(in_features=256, out_features=1, bias=True)\n",
      "  (value): Linear(in_features=256, out_features=1, bias=True)\n",
      "  (matrix_entries): Linear(in_features=256, out_features=1, bias=True)\n",
      ")\n",
      "Episode 0\tFrame 30\tAverage Score: 30.000\n",
      "Episode 100\tFrame 1355\tAverage Score: 13.255\n",
      "Episode 200\tFrame 3446\tAverage Score: 20.911\n",
      "Episode 300\tFrame 8908\tAverage Score: 54.622\n",
      "Episode 400\tFrame 18557\tAverage Score: 96.499\n",
      "Episode 495\tFrame 29919 \tAverage Score: 118.37Training time: 6.84min\n"
     ]
    }
   ],
   "source": [
    "    def run(frames=1000):\n",
    "        \"\"\"\"NAF.\n",
    "\n",
    "        Params\n",
    "        ======\n",
    "\n",
    "        \"\"\"\n",
    "        scores = []                        # list containing scores from each episode\n",
    "        scores_window = deque(maxlen=100)  # last 100 scores\n",
    "        frame = 0\n",
    "        i_episode = 0\n",
    "        state = env.reset()\n",
    "        score = 0                  \n",
    "        for frame in range(1, frames+1):\n",
    "            action = agent.act(state)\n",
    "\n",
    "            next_state, reward, done, _ = env.step(np.array([action]))\n",
    "            agent.step(state, action, reward, next_state, done, writer)\n",
    "\n",
    "            state = next_state\n",
    "            score += reward\n",
    "\n",
    "            if done:\n",
    "                scores_window.append(score)       # save most recent score\n",
    "                scores.append(score)              # save most recent score\n",
    "                writer.add_scalar(\"Reward\", score, i_episode)\n",
    "                writer.add_scalar(\"Average100\", np.mean(scores_window), i_episode)\n",
    "                print('\\rEpisode {}\\tFrame {} \\tAverage Score: {:.2f}'.format(i_episode, frame, np.mean(scores_window)), end=\"\")\n",
    "                if i_episode % 100 == 0:\n",
    "                    print('\\rEpisode {}\\tFrame {}\\tAverage Score: {:.2f}'.format(i_episode,frame, np.mean(scores_window)))\n",
    "                i_episode +=1 \n",
    "                state = env.reset()\n",
    "                score = 0              \n",
    "\n",
    "        return np.mean(scores_window)\n",
    "\n",
    "\n",
    "    if __name__ == \"__main__\":\n",
    "\n",
    "        writer = SummaryWriter(\"runs/\"+\"NAF_test_128_relu_Multinorm\")\n",
    "        seed = 0\n",
    "        BUFFER_SIZE = 100000\n",
    "        BATCH_SIZE = 128\n",
    "        GAMMA = 0.99\n",
    "        TAU = 1e-3\n",
    "        LR = 1e-3\n",
    "        UPDATE_EVERY = 1\n",
    "        NUPDATES = 1\n",
    "        device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "        print(\"Using \", device)\n",
    "\n",
    "\n",
    "        np.random.seed(seed)\n",
    "        env = gym.make(\"CartPoleConti-v0\")\n",
    "\n",
    "\n",
    "        env.seed(seed)\n",
    "        action_size     = env.action_space.shape[0]\n",
    "        state_size = env.observation_space.shape[0]\n",
    "\n",
    "        agent = DQN_Agent(state_size=state_size,    \n",
    "                            action_size=action_size,\n",
    "                            layer_size=256,\n",
    "                            BATCH_SIZE=BATCH_SIZE, \n",
    "                            BUFFER_SIZE=BUFFER_SIZE, \n",
    "                            LR=LR, \n",
    "                            TAU=TAU, \n",
    "                            GAMMA=GAMMA, \n",
    "                            UPDATE_EVERY=UPDATE_EVERY,\n",
    "                            NUPDATES=NUPDATES,\n",
    "                            device=device, \n",
    "                            seed=seed)\n",
    "\n",
    "\n",
    "\n",
    "        t0 = time.time()\n",
    "        final_average100 = run(frames = 30000)\n",
    "        t1 = time.time()\n",
    "\n",
    "        print(\"Training time: {}min\".format(round((t1-t0)/60,2)))\n",
    "        torch.save(agent.qnetwork_local.state_dict(), \"NAF_test\"+\".pth\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
